# -*- coding: utf-8 -*-
"""
Created on Wed Dec  8 11:48:14 2021

@author: zhiyinan
"""
# This module contains the casidi version of the finite difference model for the 1D Richards Equation (RE). 
# The 1D RE is expressed as C(h)dh/dt = d/dz(K(h)d/dz(h+z) where the terms have the following meanings
# C(h) -  differential water capacity/ capillary capacity
# h is -  the soil water pressure head
# K(h) -  is the unsaturated hydraulic conductivity
# z is -  the vertival coordinate directed upward

# Assumptions
# 1. The field is homogeneous (i.e same soil type) and no slope exists.
# 2. This module does not include the crop model.

# Boundary Conditions
# 1. Top BC    - Flux boundary condition [dh/dz= -1 - irrigation_amount/K(h)]
# 2. Bottom BC - Free Drainage [d(h+z)/dz= 1]

# The dependence of C, and K on h can be found in the van_GenucthenNumpy module.



# import numpy as np 
from casadi import *
from Parameters1D_casadi import *
import van_Genuchten1D_casadi as vg 


def Richards1D_casadi(head, irrigAmount, cropCoeff, refEvap):
    


    ##The spatial parameters
    totalDepth, axialNodes, axialDistance, totalNodes = spatialVariables_1D()

    capCapacity = vg.capillaryCapacity(head) 

    ##Creating an array for the flux  at each node

    flux = []
    
    # flux = np.zeros(totalNodes+1)

    

    ##Bottom boundary
    #For the free drainage boundary condition, the flux equals the hydraulic conductivity of the bottom node
    flux = vertcat(flux, -1*vg.hydraulicConductivity(head[0]))

    ##Flux for the internal nodes

    #The hydraulic conductivity at the center of the compartments
    # cntr_1 = np.arange(0, totalNodes-1)
    nodalHydCond = vg.hydraulicConductivity(head)
    approxHydCond = (nodalHydCond[1:totalNodes] + nodalHydCond[0:totalNodes-1])/2.0

    # cntr_2 = np.arange(1,totalNodes)

    flux = vertcat(flux, -approxHydCond*(((head[1:totalNodes]-head[0:totalNodes-1])/axialDistance) + 1.0))
    
    ##Top boundary condition
    flux = vertcat(flux, irrigAmount)

    
    # sinkTerms = np.zeros((1, axialNodes))
    # sinkTerms = np.zeros(axialNodes)
    
    #dhdt = (-(flux[cntr_3 + 1] - flux[cntr_3])/ axialDistance)/capCapacity - ((cropCoeff*refEvap)/totalDepth)/capCapacity +  procDisturbance
    
    # i= np.arange(0, 1)
    distance = axialDistance*DM(np.arange(0, axialNodes))
    sinkTerms = (2*(refEvap * cropCoeff)/ totalDepth)*(1-(totalDepth - distance)/totalDepth) #Weather conditions
    
    
    
    #dhdt = ((-(flux[cntr_3 + 1] - flux[cntr_3])/ axialDistance) - ((cropCoeff*refEvap)/totalDepth))/capCapacity
    
    # cntr_3 = np.arange(0, totalNodes)
    
    dhdt = ((-(flux[1:totalNodes+1] - flux[0:totalNodes])/ axialDistance) - sinkTerms)/capCapacity
    return dhdt



def volMoistureAllNodes_1D(head):
    
    totalDepth, axialNodes, axialDistance, totalNodes = spatialVariables_1D()
    volMoistureInit = vg.volumetricMoisture(head)
    cMatrix = CMatrixAllMeasurements(totalNodes)
    volMoistureFin = mtimes(cMatrix, volMoistureInit)
    return volMoistureFin


##For this model,the cMatrix is defined as follows

def cMatrix_1D(axialNodes):

    cMatrix = DM.zeros((1,axialNodes))

    for i in range(axialNodes):
        cMatrix[0,i] = (1/axialNodes)
        
    # cMatrix = np.reshape(cMatrix, (1,16))

    return cMatrix



def volMoistureSelected_1D(cMatrix, head):
    volMoistureInit = vg.volumetricMoisture(head)
    volMoistureFin = mtimes(cMatrix, volMoistureInit)
    return volMoistureFin